#!/usr/bin/python


"""

parse_gene_structure_from_vulgar.py

    --vulgar_file PATH

    reads (e.g. Exonerate) files containing one or lines in the VULGAR
    format. These should have "vulgar:" as the first letters on the line.

    If this programme is for parsing Exonerate output, remember to specify
    at least the following option:
    --showvulgar TRUE
    

"""

from optparse import OptionParser
from string import *
import sys, os
import StringIO
import string
from itertools import izip

# add self to search path for testing
if __name__ == '__main__':
    exe_path = os.path.split(os.path.abspath(sys.argv[0]))[0]
    sys.path.append(os.path.abspath(os.path.join(exe_path,"..", "python_modules")))
    myname = os.path.split(sys.argv[0])[1]
    myname = os.path.splitext(myname)[0];
else:
    myname = __name__

    
from general_util import die_error
from exon_structure import gene_exons




parser = OptionParser(version="%prog 1.0")

#
#   input files
# 
parser.add_option("-x", "--vulgar_file", dest="vulgar_file",
                  metavar="FILE", 
                  type="string",
                  help="Name and path of results file with lines in the vulgar format. "\
                        "Exonerate should have been run with the following options: "\
                        "--showvulgar TRUE --showtargetgff TRUE --showquerygff FALSE")
parser.add_option("-d", "--dna_file", dest="dna_file",
                  metavar="FILE", 
                  type="string",
                  help="Name and path of file with genomic DNA. "+
                        "This must contain either 'raw' DNA or be in the FASTA format.")
parser.add_option("-q", "--query_dna_file", dest="query_dna_file",
                  metavar="FILE", 
                  type="string",
                  help="Name and path of file with query cDNA. "+
                        "This must contain either 'raw' DNA or be in the FASTA format.")
parser.add_option("-i", "--index", dest="indices",
                  metavar="NUMBERS", 
                  type="string",
                  help="Which predictions to output (counting from zero) in a comma "+
                        "separated list.")




#
#   Output files
# 
parser.add_option("-C", "--compare_cdna_by_exons", dest="compare_cdna_by_exons_file",
                  metavar="FILE", 
                  type="string",
                  help="Name and path of file to hold aligned cDNA broken up into "+
                        "exons.")


parser.add_option("-g", "--genomic_by_exons", dest = "genomic_by_exons_file",
                  metavar="FILE",                       
                  type="string",
                  help="name and path of file for genomic DNA sequence broken up into "\
                        "introns and exons.")
parser.add_option("-e", "--exons", dest="exons_file",
                  metavar="FILE",
                  type="string",
                  help="name and path of file for exons.")
parser.add_option("-c", "--cds_in_exons", dest="cds_file",
                  metavar="FILE", 
                  type="string",
                  help="Name and path of file to hold cDNA broken up into exons.")

parser.add_option("-p", "--prot_in_exons", dest="prot_in_exons_file",
                  metavar="FILE", 
                  type="string",
                  help="Name and path of file to hold translated peptide sequence "\
                        "broken up into exons.")

parser.add_option("-P", "--peptide", dest="peptide_file",
                  metavar="FILE", 
                  type="string",
                  help="Name and path of file to hold continguous translated peptide sequence.")

parser.add_option("-l", "--lengths_and_splice", dest = "lengths_and_splice_file",
                  metavar="FILE", 
                  type="string",
                  help="Name and path of file to hold exon and intron lengths and splice_sites.")
parser.add_option("-s", "--summary", dest="summary_file",
                  metavar="FILE", 
                  type="string",
                  help="Name and path of file to total lengths of different GFFs.")
parser.add_option("-w", "--wrap_genomic", dest = "wrap_genomic",
                  action="store_true", default=False,
                  help="line wrap (main) genomic output file every 60 residues")

file_parameters = """
dna_file
vulgar_file
genomic_by_exons_file
exons_file
cds_file
prot_in_exons_file
peptide_file
lengths_and_splice_file
summary_file""".split()

#
#   everything defaults to empty string
# 
for parameter in file_parameters :
    parser.set_defaults(**{parameter: None}) 




# get help string
f =StringIO.StringIO()
parser.print_help(f)
helpstr = f.getvalue()
(options, args) = parser.parse_args()



# 
# mandatory options
# 
for parameter in ["dna_file","vulgar_file"] :
    if getattr(options, parameter) == None:
        die_error("--dna_file and --vulgar_file are mandatory. "\
                "Please specify these files.\n\n" + helpstr)

# 
# indices
# 
if options.indices != None:
    import re
    options.indices = [int(x) for x in re.split("\s+", options.indices)]


# 
# compare cDNA exon by exon
# 
if (options.compare_cdna_by_exons_file != None and
    options.query_dna_file == ""):
    die_error("Please specify the query cDNA in --query_dna_file if you wish to compare "
              "the query and target cDNAs in --compare_cdna_by_exons_file.")



    
#-----------------------------------------------------------------------------------------
#
#       parses exon structure from "vulgar" line in exonerate results file
# 
# 
def parse_gene_structure_from_vulgar_line(vulgar_line):
    data_fields = vulgar_line.split()
    
    # split off header fields
    (vulgar_name, 
     query_id, query_beg, query_end, query_strand,
     target_id, target_beg, target_end, target_strand,
     score)                                              = data_fields[0:10]
    data_fields = data_fields[10:]
    
    # e.g.
    # vulgar: 
    # prot.fa       0       3162    . 
    # canFam2_dna   10000   1531527 + 
    # 12025

    new_gene = gene_exons(target_id, target_strand)

    #
    #   current code limitations in terms of strand
    # 
    if query_strand != '.':
        raise "Error! Query strand = %s This code assumes "\
                "that we are building genes off proteins (\n%s)" % (query_strand, line)
    if target_strand != '+':
        raise "Error! This code assumes +ve strand"

    curr_exonic_pos = 0
    curr_target_pos = int(target_beg)
    exon_start      = curr_target_pos
    exon_finish     = curr_target_pos
    exonic_coding_regions = list()

    intronic_mode = False

    for what, qq, tt in izip(data_fields[::3], data_fields[1::3], data_fields[2::3]):
        
        # e.g.
        # what  qq  tt
        # ----  --  ----
        # M     22  66      match               # add to coding
        # G     0   3       gap                 # ignore gaps: this is a legit indel
        # M     227 681     match               
        # S     0   1       split codon         # add to coding
        # 5     0   2       5' splice site      # add to intron
        # I     0   3462    intron              # add to intron
        # 3     0   2       3' splice site      # add to intron 
        # S     1   2       split codon         # add to coding
        # M     37  111     new exon            # new exon
        
        
        qq, tt = (int(qq), int(tt))
        # M match
        # S split codon
        # G gap (in-frame indel)
        if what in 'MSG':
            # if coming out of an intron, add the accumalated data as an exon
            if intronic_mode:
                new_gene.add_exon(exon_start, exon_finish, exonic_coding_regions, '')
                #if not len(new_gene.exons):
                #    print new_gene.first_exon.coding_regions
                #else:
                #    print new_gene.exons[-1].coding_regions
                intronic_mode = False
                exon_start = curr_target_pos
                exonic_coding_regions = []
            exonic_coding_regions.append([curr_exonic_pos, curr_exonic_pos + tt])
            curr_exonic_pos += tt
                
        # splice site
        elif what in "5I3":
            if not intronic_mode:
                intronic_mode = True
                exon_finish = curr_target_pos
                
            # reset exonic pos
            curr_exonic_pos = 0
            
        #F frame shift
        elif what in "F":
            curr_exonic_pos += tt 
        #N    Non-equivalenced region
        else:
            raise "Don't understand Non-equivalenced region (%s, %d, %d)" % (
                                                                what, qq, tt)

        curr_target_pos += tt
            
    else:
        new_gene.finish_init()
        return new_gene

    raise "Logic error"

#
#       parses exon structure from "vulgar" line in exonerate results file
# 
# 
#-----------------------------------------------------------------------------------------
            

                        
                    
        



#
#  read dna
# 

from fasta_sequences import fasta_sequences

#  target genomic dna
target_dna = fasta_sequences()
target_dna.read_dna_from_file (options.dna_file)

#  query cDNA
query_dna = fasta_sequences()
if options.query_dna_file:
    query_dna.read_dna_from_file (options.query_dna_file)

    


#
#  read exonerate results
# 
vulgar_file = open(options.vulgar_file)
exonerate_genes = list()
for line in vulgar_file:
    if line[0:7] == "vulgar:":
        exonerate_genes.append(parse_gene_structure_from_vulgar_line(line))




#
#   open output files
# 
def open_or_stdout(file_name):
    if file_name == None: 
        return None
    if file_name == '-':
        return sys.stdout
    else:
        return open(file_name,  "w")

exons_file                  = open_or_stdout(options.exons_file         )
prot_in_exons_file          = open_or_stdout(options.prot_in_exons_file )
peptide_file                = open_or_stdout(options.peptide_file )
lengths_and_splice_file     = open_or_stdout(options.lengths_and_splice_file)
summary_file                = open_or_stdout(options.summary_file)
genomic_by_exons_file       = open_or_stdout(options.genomic_by_exons_file)
compare_cdna_by_exons_file  = open_or_stdout(options.compare_cdna_by_exons_file)
cds_file                    = open_or_stdout(options.cds_file)






        
#
# if not specified, go through all genes
# 
sys.stderr.write("\t%d gene structures predicted\n" % (len(exonerate_genes)))
if options.indices == None:
    options.indices = range(len(exonerate_genes))

#
#   helper function to write out exons/introns
# 
def write_intron_exons (fileH, gene, iter):
    fileH.write("#" + gene.get_summary() + "\n")
    for description, seq in iter (target_dna[gene.contig], 500):
        fileH.write( ">" + description + "\n")
        # wrap lines
        if options.wrap_genomic:
            fileH.write(line_wrap(seq, 60) + "\n")
        else:
            fileH.write(seq + "\n")

def write_exons (fileH, gene, iter):
    fileH.write("#" + gene.get_summary() + "\n")
    for description, seq in izip(gene.iter_descriptions(), iter):
        fileH.write( ">" + description + "\n")
        # wrap lines
        if options.wrap_genomic:
            fileH.write(line_wrap(seq, 60) + "\n")
        else:
            fileH.write(seq + "\n")
    
    
for x in options.indices:

    # ignore out of range
    if x >= len(exonerate_genes):
        sys.stderr.write("\tPredicted gene index #%d is out of range and will be ignored.\n" %
                         x)
        continue;

    gene = exonerate_genes[x]

    # 
    # splice sites and lengths
    # 
    if lengths_and_splice_file:
        line_format = "%-9s " * 7  + "\n"
        lengths_and_splice_file.write(line_format % gene.tabulated_summary_columns())
        for cols in gene.tabulated_summary(target_dna[gene.contig]):
            lengths_and_splice_file.write(line_format % tuple(cols))

    #
    #   summary
    # 
    if summary_file:
        summary_file.write(gene.get_summary() + "\n")
    
    #
    #   genomic intron exon
    # 
    if genomic_by_exons_file:
        write_intron_exons(genomic_by_exons_file, gene, gene.iter_intron_exons)
            
    #
    #   all exons
    # 
    if exons_file:
        write_exons(exons_file, gene, gene.iter_exons(target_dna[gene.contig]))
            


    #
    #   translations by exons
    # 
    if prot_in_exons_file:
        write_exons(prot_in_exons_file, gene, gene.iter_translation(target_dna[gene.contig], True))


    if peptide_file:
        peptide_file.write(">peptide " + gene.get_summary() + "\n")
        for seq in gene.iter_translation(target_dna[gene.contig], False):
            # wrap lines
            if options.wrap_genomic:
                peptide_file.write(line_wrap(seq, 60) + "\n")
            else:
                peptide_file.write(seq + "\n")


    #
    #   CDS by exons
    # 
    if cds_file:
        write_exons(cds_file, gene, gene.iter_cds(target_dna[gene.contig]))
                        



    #if compare_cdna_by_exons_file:
    #    gene.print_matching_exon_sequence(target_dna[gene.contig], query_dna, compare_cdna_by_exons_file)








#    def print_matching_exon_sequence (self, dna_target, dna_query, file):
#        """
#        print_matching_exon_sequence
#
#            prints out exons sequence matching query to target
#            in fasta format
#            
#        """
#        from translate_by_codon import translate
#        
#        self.print_summary (file)
#        dna_query = dna_query[0]
#        dna_target = dna_target[self.contig]
#
#        differences_per_exon = list()
#        matching_exon_lengths = list()
#        curr_frame = 0
#        for exon_index, coordinates in enumerate(self.gaps):
#            exon_start, exon_finish, exon_comment = self.exons[exon_index]
#            # write out exon
#            file.write(">exon %d-%d len = %d %s\n" % (exon_start, exon_finish, 
#                                                      exon_finish - exon_start, 
#                                                      exon_comment))
#            query_seq = ''
#            target_seq = ''
#            for q_pos, t_pos, length in coordinates:
#                query_seq += dna_query[q_pos:q_pos + length]
#                target_seq += dna_target[exon_start + t_pos: exon_start + t_pos + length]
#
#            file.write(query_seq + "\n")
#            file.write(target_seq + "\n")
#            file.write(translate(query_seq [(3-curr_frame) % 3:]) + "\n")
#            file.write(translate(target_seq[(3-curr_frame) % 3:]) + "\n")
#
#            curr_frame = (len(query_seq) + curr_frame) % 3
#            
#
#            #
#            #   calculate differences at each position for this exon
#            # 
#            differences_per_exon.append([0,0,0])
#            matching_exon_lengths.append(len(query_seq))
#            for offset in range(3):
#                for s1, s2 in izip(query_seq[offset::3], target_seq[offset::3]):
#                    if s1 != s2:
#                        differences_per_exon[-1][offset] += 1
#
#        for (pos1, pos2, pos3), length in izip(differences_per_exon, matching_exon_lengths):
#            file.write("%5d\t%5d\t%5d\t%5d\t%.2f\t%.2f\t%.2f\n" %(length, pos1, pos2, pos3, 
#                                                    round(pos1 * 100.0 / length, 2),
#                                                    round(pos2 * 100.0 / length, 2),
#                                                    round(pos3 * 100.0 / length, 2)))

